"""
嵌入API封装
支持DeepSeek和通义千问作为备用方案
"""

import os
import logging
import requests
from typing import List, Dict, Any, Optional
import openai
from dotenv import load_dotenv

# 加载环境变量
load_dotenv()

logger = logging.getLogger(__name__)

class Embeddings:
    """嵌入API封装类，支持多后端"""
    
    def __init__(self, backend: str = "qwen"):
        """初始化嵌入客户端
        
        Args:
            backend: 嵌入后端，支持"deepseek"或"qwen"
        """
        self.backend = backend
        
        if backend == "deepseek":
            self._init_deepseek()
        elif backend == "qwen":
            self._init_qwen()
        else:
            raise ValueError(f"不支持的嵌入后端: {backend}")
    
    def _init_deepseek(self):
        """初始化DeepSeek客户端"""
        self.api_key = os.getenv('DEEPSEEK_API_KEY')
        self.base_url = os.getenv('DEEPSEEK_BASE_URL', 'https://api.deepseek.com/v1')
        
        if not self.api_key:
            raise ValueError("DEEPSEEK_API_KEY 环境变量未设置")
        
        self.client = openai.OpenAI(
            api_key=self.api_key,
            base_url=self.base_url
        )
        
        # 测试API连接
        try:
            self.client.models.list()
            logger.info("DeepSeek API 连接测试成功")
        except Exception as e:
            logger.warning(f"DeepSeek API 连接失败: {str(e)}")
            raise ValueError(f"无法连接到DeepSeek API: {str(e)}")
    
    def _init_qwen(self):
        """初始化通义千问客户端"""
        import dashscope
        self.api_key = os.getenv('QWEN_API_KEY')
        
        if not self.api_key:
            raise ValueError("QWEN_API_KEY 环境变量未设置")
        
        dashscope.api_key = self.api_key
        logger.info("通义千问嵌入初始化成功")
    
    def get_backend(self) -> str:
        """获取当前使用的嵌入后端"""
        return self.backend
    
    def embed_text(self, text: str) -> List[float]:
        """
        将单个文本转换为向量
        
        Args:
            text: 要向量化的文本
            
        Returns:
            文本的向量表示
        """
        try:
            if self.backend == "deepseek":
                response = self.client.embeddings.create(
                    model="text-embedding-3-small",
                    input=text
                )
                return response.data[0].embedding
            elif self.backend == "qwen":
                return self._embed_with_qwen(text)
        except Exception as e:
            logger.error(f"文本向量化失败: {str(e)}")
            raise
    
    def _embed_with_qwen(self, text: str) -> List[float]:
        """使用通义千问进行嵌入"""
        from dashscope import TextEmbedding
        resp = TextEmbedding.call(
            model="text-embedding-v1",
            input=text
        )
        if resp.status_code != 200:
            raise ValueError(f"通义千问API请求失败: {resp}")
        return resp.output["embeddings"][0]["embedding"]
    
    def embed_texts(self, texts: List[str]) -> List[List[float]]:
        """
        批量将文本转换为向量
        
        Args:
            texts: 要向量化的文本列表
            
        Returns:
            文本向量列表
        """
        try:
            if self.backend == "deepseek":
                response = self.client.embeddings.create(
                    model="text-embedding-3-small",
                    input=texts
                )
                return [data.embedding for data in response.data]
            elif self.backend == "qwen":
                return self._embed_batch_with_qwen(texts)
        except Exception as e:
            logger.error(f"批量文本向量化失败: {str(e)}")
            raise
    
    def _embed_batch_with_qwen(self, texts: List[str]) -> List[List[float]]:
        """使用通义千问进行批量嵌入"""
        from dashscope import TextEmbedding
        resp = TextEmbedding.call(
            model="text-embedding-v1",
            input=texts
        )
        if resp.status_code != 200:
            raise ValueError(f"通义千问API请求失败: {resp}")
        return [item["embedding"] for item in resp.output["embeddings"]]
    
    def get_embedding_dimension(self) -> int:
        """获取嵌入向量的维度"""
        # DeepSeek embedding 模型的维度
        return 1536
    
    def similarity(self, vec1: List[float], vec2: List[float]) -> float:
        """
        计算两个向量的余弦相似度
        
        Args:
            vec1: 第一个向量
            vec2: 第二个向量
            
        Returns:
            相似度分数 (0-1)
        """
        import numpy as np
        
        vec1 = np.array(vec1)
        vec2 = np.array(vec2)
        
        # 计算余弦相似度
        dot_product = np.dot(vec1, vec2)
        norm1 = np.linalg.norm(vec1)
        norm2 = np.linalg.norm(vec2)
        
        if norm1 == 0 or norm2 == 0:
            return 0.0
        
        return dot_product / (norm1 * norm2)
